import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

# 线性回归公式 y = w1*x1 + w2*x2 + wn*xn + b
# 其中 w1, w2, wn 为权重，b 为偏置
# y是预测值，x是输入特征

# 1、准备数据
X = torch.randn(100, 2)  # 随机生成100个样本，每个样本有2个特征
t_w = torch.tensor([2.0, 3.0])  # 真实权重
t_b = 4.0  # 真实偏置
Y = X @ t_w + t_b + torch.randn(100) * 0.1  # 加入一些噪声

# 打印数据
print("X:", X[:5])  # 打印前5个样本
print("Y:", Y[:5])  # 打印前5个标签


# 2、定义线性回归模型
class LinearRegressionModel(nn.Module):
    def __init__(self):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(2, 1)  # 输入特征为2，输出特征为1

    def forward(self, x):
        return self.linear(x)  # 前向传播
model = LinearRegressionModel()  # 创建模型实例

# 3、定义损失函数和优化器
criterion = nn.MSELoss()  # 均方误差损失函数
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 随机梯度下降优化器

# 4、训练模型
epochs = 1000  # 训练轮数
for epoch in range(epochs):
    model.train()  # 设置模型为训练模式

    # 前向传播
    outputs = model(X)  # 前向传播
    loss = criterion(outputs, Y.view(-1, 1))  # 计算损失

    # 反向传播
    optimizer.zero_grad()
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

    if (epoch + 1) % 100 == 0:  # 每100轮打印一次损失
        print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}')  # 打印损失

# 5、评估模型、查看训练后的权重和偏置等参数
print(f'w:{model.linear.weight.data.numpy()}')
print(f'bias"{model.linear.bias.data.numpy()}')

# 在新数据上做预测
with torch.no_grad():  # 评估时不需要计算梯度
    predictions = model(X)

# 可视化预测与实际值
plt.scatter(X[:, 0], Y, color='blue', label='True values')
plt.scatter(X[:, 0], predictions, color='red', label='Predictions')
plt.legend()
plt.show()